﻿#include "dmm.h"
#include "vk_mem_alloc.h"
#include <cstdint>
#include <stdexcept>
#include <cassert>
#include <string.h>

#define STRINGIZE(x) STRINGIZE2(x)
#define STRINGIZE2(x) #x
#define LINE_STRING STRINGIZE(__LINE__)
#define TEST(expr)  do { if(!(expr)) { \
        assert(0 && #expr); \
        throw std::runtime_error(__FILE__ "(" LINE_STRING "): ( " #expr " ) == false"); \
    } } while(false)
#define ERR_GUARD_VULKAN(expr)  do { if((expr) < 0) { \
        assert(0 && #expr); \
        throw std::runtime_error(__FILE__ "(" LINE_STRING "): VkResult( " #expr " ) < 0"); \
    } } while(false)

namespace dmm
{
	/* CPU Device Memory Functions */
	static void CPUGetPhysicalDeviceMemoryProperties(VkPhysicalDevice physicalDevice, VkPhysicalDeviceMemoryProperties* pMemoryProperties)
	{
		pMemoryProperties->memoryTypeCount = 3;
		//支持的memory property
		pMemoryProperties->memoryTypes[0].propertyFlags = 0x7;
		pMemoryProperties->memoryTypes[0].heapIndex = 0;
		pMemoryProperties->memoryTypes[1].propertyFlags = 0xf;
		pMemoryProperties->memoryTypes[1].heapIndex = 0;
		pMemoryProperties->memoryTypes[2].propertyFlags = 0x1;
		pMemoryProperties->memoryTypes[2].heapIndex = 0;

		pMemoryProperties->memoryHeapCount = 1;
		//总的heap size。该值会影响VmaAllocator_T创建BlockVector时对preferredBlockSize的计算
		//若preferredBlockSize比较小，那么在申请内存时更容易走dedicatedPreferred分支
		//这个值只是一个“指示”值，并不能限制应用程序一共申请多大的内存。
		//即使把这个值设置为0，也不影响内存申请。若想限制内存申请总量，需使用VmaAllocatorCreateInfo::pHeapSizeLimit
		pMemoryProperties->memoryHeaps[0].size = 8ULL << 30;	//8G
		pMemoryProperties->memoryHeaps[0].flags = 1;
	}
	//简化的GetPhysicalDeviceProperties
	static void CPUGetPhysicalDeviceProperties(VkPhysicalDevice physicalDevice, VkPhysicalDeviceProperties* pProperties)
	{
		pProperties->apiVersion = VK_API_VERSION_1_2;
		strcpy(pProperties->deviceName, "VIRTUAL_GPU");
		pProperties->deviceType = VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU;
		pProperties->limits.nonCoherentAtomSize = 1;
		pProperties->limits.bufferImageGranularity = 1;
		pProperties->limits.maxMemoryAllocationCount = 4096;
	}
	static VkResult CPUAllocateMemory(VkDevice device, const VkMemoryAllocateInfo* pAllocateInfo, const VkAllocationCallbacks* pAllocator, VkDeviceMemory* pMemory)
	{
		void* p = malloc(pAllocateInfo->allocationSize);
		if (nullptr == pMemory)
		{
			*pMemory = (VkDeviceMemory)nullptr;
			return VK_ERROR_OUT_OF_DEVICE_MEMORY;
		}
		else
		{
			*pMemory = (VkDeviceMemory)p;
			return VK_SUCCESS;
		}
	}
	static void CPUFreeMemory(VkDevice device, VkDeviceMemory memory, const VkAllocationCallbacks* pAllocator)
	{
		free((void*)memory);
	}
	static VkResult CPUMapMemory(VkDevice device, VkDeviceMemory memory, VkDeviceSize offset, VkDeviceSize size, VkMemoryMapFlags flags, void** ppData)
	{
		*ppData = (void*)memory;
		return VK_SUCCESS;
	}
	static void CPUUnmapMemory(VkDevice device, VkDeviceMemory memory)
	{
		return;
	}
	static VkResult CPUFlushMappedMemoryRanges(VkDevice device, uint32_t memoryRangeCount, const VkMappedMemoryRange* pMemoryRanges)
	{
		return VK_SUCCESS;
	}
	static VkResult CPUInvalidateMappedMemoryRanges(VkDevice device, uint32_t memoryRangeCount, const VkMappedMemoryRange* pMemoryRanges)
	{
		return VK_SUCCESS;
	}
	VmaVulkanFunctions CPUDeviceFunctions = {
		&CPUGetPhysicalDeviceProperties,
		&CPUGetPhysicalDeviceMemoryProperties,
		&CPUAllocateMemory,
		&CPUFreeMemory,
		&CPUMapMemory,
		&CPUUnmapMemory,
		&CPUFlushMappedMemoryRanges,
		&CPUInvalidateMappedMemoryRanges
	};

	VmaAllocator g_hAllocator = nullptr;

	static size_t dmmAlignUp(size_t val, size_t alignment)
	{
		assert((alignment & (alignment - 1)) == 0);
		return (val + alignment - 1) & ~(alignment - 1);
	}

	void dmmInit(void)
	{
		VmaAllocatorCreateInfo allocatorInfo = {};
		allocatorInfo.vulkanApiVersion = VK_API_VERSION_1_2;
		allocatorInfo.pVulkanFunctions = &CPUDeviceFunctions;	//必须提供
		ERR_GUARD_VULKAN(vmaCreateAllocator(&allocatorInfo, &g_hAllocator));
	}

	void* dmmAlloc(size_t len, size_t alignment)
	{
#if 0
		VmaAllocationCreateInfo allocCreateInfo = {};
		allocCreateInfo.flags = VMA_ALLOCATION_CREATE_MAPPED_BIT        //总是需要用户地址
			| VMA_ALLOCATION_CREATE_WITHIN_BUDGET_BIT;

		size_t additionInfoSize = sizeof(VmaAllocation);
		VkMemoryRequirements memReq = {};
		memReq.size = 200 << 20;
		memReq.alignment = 0;									//对底层的align不做要求，待结果返回后自己做align
		memReq.memoryTypeBits = UINT32_MAX;						//可接受所有的memory type

		VmaAllocation alloc = VK_NULL_HANDLE;
		VmaAllocationInfo allocationInfo;
		VkResult res = vmaAllocateMemory(g_hAllocator, &memReq, &allocCreateInfo, &alloc, &allocationInfo);
		size_t adata = (size_t)allocationInfo.pMappedData;

		//VmaPoolCreateInfo poolCreateInfo = {};
		//VmaPool pool;
		//vmaCreatePool(g_hAllocator, &poolCreateInfo, &pool);
		//allocCreateInfo.pool = pool;
		//res = vmaAllocateMemory(g_hAllocator, &memReq, &allocCreateInfo, &alloc, &allocationInfo);
		//adata = (size_t)allocationInfo.pMappedData;

		return nullptr;

#else
		VmaAllocationCreateInfo allocCreateInfo = {};
		allocCreateInfo.flags = VMA_ALLOCATION_CREATE_MAPPED_BIT;	//总是需要用户地址

		size_t additionInfoSize = sizeof(VmaAllocation);
		VkMemoryRequirements memReq = {};
		memReq.size = len + alignment + additionInfoSize;
		memReq.alignment = 0;									//对底层的align不做要求，待结果返回后自己做align
		memReq.memoryTypeBits = UINT32_MAX;						//可接受所有的memory type

		VmaAllocation alloc = VK_NULL_HANDLE;
		VmaAllocationInfo allocationInfo;
		VkResult res = vmaAllocateMemory(g_hAllocator, &memReq, &allocCreateInfo, &alloc, &allocationInfo);
		if (VK_SUCCESS != res)
		{
			return nullptr;
		}
		else
		{
			size_t adata = (size_t)allocationInfo.pMappedData;
			size_t udata = dmmAlignUp(adata + additionInfoSize, alignment);

			void** pudata = (void**)udata;
			pudata[-1] = (void*)alloc;

			return (void*)udata;
		}
#endif
	}

	void dmmFree(void* addr)
	{
		void** pudata = (void**)addr;
		VmaAllocation alloc = (VmaAllocation)pudata[-1];
		vmaFreeMemory(g_hAllocator, alloc);
	}

	void dmmDeinit(void)
	{
		vmaDestroyAllocator(g_hAllocator);
	}

	std::string dmmGetMemUsageInfo(void)
	{
		char* statsStr = nullptr;
        vmaBuildStatsString(g_hAllocator, &statsStr, VK_TRUE);
		std::string str(statsStr);
        vmaFreeStatsString(g_hAllocator, statsStr);

		return str;
	}
}
